{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Experiment with ALFWorld environment"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Indicate the LLM version that you want to use\n",
    "# 'gpt-3.5' / 'gpt-3'\n",
    "GPT_MODEL = 'gpt-3'\n",
    "\n",
    "# max number of trials for each task\n",
    "NUM_TRY = 6\n",
    "# after NUM_TRY_RESET trials, the agent will try to start from step 1.\n",
    "NUM_TRY_RESET = 3\n",
    "\n",
    "# By default, the max number of tokens is 1400. In some tasks, the context limit of the language models is exceeded. Try to change the number of tokens in these cases.\n",
    "MAX_TOKENS = 1400\n",
    "\n",
    "prefixes = {\n",
    "    'pick_and_place': 'put',\n",
    "    'pick_clean_then_place': 'clean',\n",
    "    'pick_heat_then_place': 'heat',\n",
    "    'pick_cool_then_place': 'cool',\n",
    "    'look_at_obj': 'examine',\n",
    "    'pick_two_obj': 'puttwo'\n",
    "}\n",
    "\n",
    "# save to files\n",
    "LOG_FILE = 'interaction_log.txt'\n",
    "PR_FILE = 'Prompt_Response.txt'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 539,
     "status": "ok",
     "timestamp": 1668226462186,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": 0
    },
    "id": "DxYJWBWWIswz",
    "outputId": "6b5123dd-d4f5-43e8-b8f4-3be029d8594c",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "import re\n",
    "import yaml\n",
    "import alfworld\n",
    "import alfworld.agents.environment\n",
    "from datetime import datetime\n",
    "import time\n",
    "import sys\n",
    "import io\n",
    "import traceback\n",
    " \n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]\n",
    "\n",
    "with open('base_config.yaml') as reader:\n",
    "    config = yaml.safe_load(reader)\n",
    "    \n",
    "split = \"eval_out_of_distribution\"\n",
    "\n",
    "env = getattr(alfworld.agents.environment, config[\"env\"][\"type\"])(config, train_eval=split)\n",
    "env = env.init_env(batch_size=1)\n",
    "\n",
    "def process_ob(ob):\n",
    "    if ob.startswith('You arrive at loc '):\n",
    "        ob = ob[ob.find('. ')+2:]    \n",
    "    return ob\n",
    "\n",
    "def start_episode(env):\n",
    "    ob, info = env.reset()\n",
    "    ob = '\\n'.join(ob[0].split('\\n\\n')[1:])\n",
    "    ob = process_ob(ob)\n",
    "    name = '/'.join(info['extra.gamefile'][0].split('/')[-3:-1])\n",
    "    return ob, name\n",
    "\n",
    "def interact_with_env(action):\n",
    "    observation, reward, done, info = env.step([action])\n",
    "    observation, reward, done = process_ob(observation[0]), info['won'][0], done[0]\n",
    "    return observation, reward, done\n",
    "\n",
    "def extract_receptacles_and_task(text):\n",
    "    # Extract list of receptacles\n",
    "    receptacles_pattern = r\"\\b\\w+ \\d+\\b\"\n",
    "    receptacles = re.findall(receptacles_pattern, text)\n",
    "\n",
    "    # Extract the task\n",
    "    task_pattern = r\"Your task is to:.*\"\n",
    "    task = re.search(task_pattern, text).group(0)\n",
    "\n",
    "    return receptacles, task\n",
    "\n",
    "def extract_answers(text, markers):\n",
    "    answers = []\n",
    "    for i, marker in enumerate(markers):\n",
    "        if i < len(markers) - 1:\n",
    "            next_marker = markers[i + 1]\n",
    "            pattern = fr\"{re.escape(marker)}\\s*([\\s\\S]*?)\\n*{re.escape(next_marker)}\"\n",
    "        else:\n",
    "            pattern = fr\"{re.escape(marker)}\\s*([\\s\\S]*)\"\n",
    "        \n",
    "        answer = re.search(pattern, text)\n",
    "        if answer:\n",
    "            answer = answer.group(1).strip()\n",
    "            answers.append(answer)\n",
    "        else:\n",
    "            answers.append('Not found')\n",
    "        \n",
    "    return answers\n",
    "\n",
    "def save_to_file(directory, filename, content):\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "        \n",
    "    file_path = os.path.join(directory, filename)\n",
    "    with open(filename, 'a') as file:\n",
    "        file.write(content + '\\n')\n",
    "\n",
    "def clear_file_content(filename):\n",
    "    with open(filename, 'w') as file:\n",
    "        pass\n",
    "\n",
    "def dict_to_traj(d):\n",
    "    traj = ''\n",
    "    for i, a in enumerate(d['actions']):\n",
    "        traj += f'Act {i}: {a}\\nObs {i}: {d[\"observations\"][i]}\\n'\n",
    "    return traj\n",
    "\n",
    "def get_line_starting_with(lines, start):\n",
    "    for line in lines.split('\\n'):\n",
    "        if line.strip().startswith(start):\n",
    "            return line\n",
    "    return None\n",
    "\n",
    "def get_error_step(error_message):\n",
    "    pattern = r'\\[Step (\\d+)\\]'\n",
    "    match = re.search(pattern, error_message)\n",
    "    if match:\n",
    "        return int(match.group(1))\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "def get_first_digit(input_string):\n",
    "    pattern = r'\\d'\n",
    "    match = re.search(pattern, input_string)\n",
    "    if match:\n",
    "        return int(match.group())\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "datetime = datetime.now().strftime(\"%m%d-%H%M%S\")\n",
    "save_to_folder = f'./r_{GPT_MODEL}_{datetime}/'\n",
    "PR_FILE =  save_to_folder + PR_FILE\n",
    "LOG_FILE = save_to_folder + LOG_FILE\n",
    "\n",
    "save_to_file(save_to_folder, PR_FILE, 'Start: ' + str(datetime) + '\\n')\n",
    "save_to_file(save_to_folder, LOG_FILE, 'Start: ' + str(datetime) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# this function captures the assertion error message and stores the local variables inside the solution at the breakpoint.\n",
    "def capture_output(func, agent, step=1):\n",
    "    # Store the original standard output and standard error\n",
    "    original_stdout = sys.stdout\n",
    "    original_stderr = sys.stderr\n",
    "\n",
    "    # Redirect the standard output and error to in-memory file-like objects\n",
    "    temp_stdout = io.StringIO()\n",
    "    temp_stderr = io.StringIO()\n",
    "    sys.stdout = temp_stdout\n",
    "    sys.stderr = temp_stderr\n",
    "\n",
    "    checkpoint = None\n",
    "    # Run the function and capture exceptions\n",
    "    try:\n",
    "        func(agent, start_from=step)\n",
    "    except Exception as e:\n",
    "        traceback.print_exc()\n",
    "        checkpoint = sys.exc_info()[2].tb_next.tb_frame.f_locals\n",
    "\n",
    "    # Restore the original standard output and error\n",
    "    sys.stdout = original_stdout\n",
    "    sys.stderr = original_stderr\n",
    "\n",
    "    # Get the output and error messages as strings\n",
    "    output_string = temp_stdout.getvalue()\n",
    "    error_string = temp_stderr.getvalue()\n",
    "\n",
    "    print(output_string)\n",
    "    print(error_string)\n",
    "    return error_string, checkpoint, output_string + error_string\n",
    "\n",
    "# this function modifies the header of the solution to load the variables from the breakpoint\n",
    "def modify_header(checkpoint):\n",
    "    if not checkpoint:\n",
    "        return 'def solution(agent, start_from=1):'\n",
    "    load_checkpoint = ''\n",
    "    skip_vars = ['agent', 'start_from']\n",
    "    for k,v in checkpoint.items():\n",
    "        if k not in skip_vars:\n",
    "            if type(v) == str:\n",
    "                load_checkpoint += f', {k}=\"{v}\"'\n",
    "            else:\n",
    "                load_checkpoint += f', {k}={v}'\n",
    "    header = f'def solution(agent, start_from{load_checkpoint}):'\n",
    "    return header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def ask(prompt):\n",
    "    prompt_chat = [\n",
    "            {\"role\": \"user\", \"content\": prompt.strip()},\n",
    "        ]\n",
    "    cnt = 0\n",
    "    while True:\n",
    "        try:\n",
    "            if GPT_MODEL == 'gpt-3.5':\n",
    "                response = openai.ChatCompletion.create(\n",
    "                    model=\"gpt-3.5-turbo\",\n",
    "                    messages=prompt_chat,\n",
    "                    temperature=0,\n",
    "                    max_tokens=MAX_TOKENS,\n",
    "                    top_p=1,\n",
    "                    frequency_penalty=0.0,\n",
    "                    presence_penalty=0.0,\n",
    "                    )\n",
    "                answer = response[\"choices\"][0][\"message\"]['content'].strip()\n",
    "                save_to_file(save_to_folder,\n",
    "                             PR_FILE, \n",
    "                            'Prompt: \\n' +\n",
    "                            prompt +\n",
    "                            '\\nResponse: \\n' +\n",
    "                            answer + '\\n' + '='*20 + '\\n')\n",
    "        \n",
    "                return answer\n",
    "            elif GPT_MODEL == 'gpt-3':\n",
    "                response = openai.Completion.create(\n",
    "                    model=\"text-davinci-002\",\n",
    "                    prompt=prompt,\n",
    "                    temperature=0,\n",
    "                    max_tokens=MAX_TOKENS,\n",
    "                    top_p=1,\n",
    "                    frequency_penalty=0.0,\n",
    "                    presence_penalty=0.0,\n",
    "                    )\n",
    "                answer = response[\"choices\"][0][\"text\"].strip()\n",
    "                save_to_file(save_to_folder,\n",
    "                             PR_FILE,\n",
    "                            'Prompt: \\n' +\n",
    "                            prompt +\n",
    "                            '\\nResponse: \\n' +\n",
    "                            answer + '\\n' + '='*20 + '\\n')\n",
    "                return answer\n",
    "            else:\n",
    "                raise Exception('Wrong GPT_MODEL')\n",
    "        except openai.error.RateLimitError as e:\n",
    "            retry_after = 3\n",
    "            print(f\"Rate limit error: {e}. Retrying in {retry_after} seconds.\")\n",
    "            time.sleep(retry_after)\n",
    "        except openai.error.InvalidRequestError as e:\n",
    "            # try to eliminate some parts of the prompt to reduce the number of tokens\n",
    "            eliminate_context = \"# for example: You have a list of receptacles, and you want to sort them by the likelihood of a soapbar appearing in them. You can do this by asking the assistant:\\nreceptacles = ['countertop 1', 'garbagecan 1', 'sinkbasin 2', 'sinkbasin 1', 'toilet 1', 'toiletpaperhanger 1', 'towelholder 1']\\nanswer = ask(f'Sort the list of receptacles, starting from the one a soapbar is most likely to appear: {receptacles}. You should return a Python list.')\\n# answer = ['sinkbasin 1', 'sinkbasin 2', 'countertop 1', 'towelholder 1', 'toiletpaperhanger 1', 'garbagecan 1', 'toilet 1']\"\n",
    "            prompt = prompt.replace(eliminate_context, '')\n",
    "            prompt_chat = [\n",
    "                {\"role\": \"user\", \"content\": prompt.strip()},\n",
    "            ]\n",
    "            print(f\"Exceed max: {e}.\")\n",
    "            cnt += 1\n",
    "            if cnt > 3:\n",
    "                return 'Exceed max limit. Tried 3 times. Skip this one.'\n",
    "        except openai.error.APIError:\n",
    "            cnt += 1\n",
    "            if cnt > 3:\n",
    "                return 'APIError. Tried 3 times. Skip this one.'\n",
    "        except Exception as e:\n",
    "            print(f\"An unexpected error occurred: {e}\")\n",
    "            raise"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Define the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from ast import literal_eval\n",
    "if GPT_MODEL == 'gpt-3.5':\n",
    "    from prompt import get_solution_prompt, code_check_prompt, feedback_fix_prompt, simple_example, puttwo_example, examine_example, clean_example, heat_example, cool_example, get_start_from_prompt\n",
    "else:\n",
    "    from prompt_GPT3 import get_solution_prompt, code_check_prompt, feedback_fix_prompt, simple_example, puttwo_example, examine_example, clean_example, heat_example, cool_example, get_start_from_prompt\n",
    "\n",
    "# Agent class represents the state of the agent, including its location,\n",
    "# what it's holding, and the states of receptacles and objects in the environment,\n",
    "# as well as the actions it can take.\n",
    "class Agent:\n",
    "    def __init__(self, receptacles):\n",
    "        self.location = None\n",
    "        self.holding = None\n",
    "        self.receptacles = receptacles\n",
    "        self.interaction_history = {'actions': [], 'observations': []}\n",
    "        self.is_success = False\n",
    "\n",
    "    # Note down the history of interactions with the environment\n",
    "    def add_to_history(self, action, observation):\n",
    "        self.interaction_history['actions'].append(action)\n",
    "        self.interaction_history['observations'].append(observation)\n",
    "        \n",
    "    # Get an observation from the environment after performing an action, and add it to the history\n",
    "    def observation(self, action):\n",
    "        observation, reward, done = interact_with_env(action)\n",
    "        self.add_to_history(action, observation)\n",
    "        print(f'Act: {action}\\nObs: {observation}')\n",
    "        if done:\n",
    "            self.is_success = reward\n",
    "            print('Done. Success:', reward)\n",
    "        return observation\n",
    "\n",
    "    # Go to a receptacle and update the agent's location. It returns an observation in natural language.\n",
    "    # For example, 'On the countertop 1, you see a candle 1, a cloth 2, and a soapbar 1.' = goto('countertop 1')\n",
    "    def goto(self, receptacle):\n",
    "        self.location = receptacle\n",
    "        return self.observation(f'go to {receptacle}')\n",
    "\n",
    "    # Take an object from a receptacle if the agent is not holding anything. It returns an observation in natural language.\n",
    "    # For example, 'You pick up the soapbar 1 from the towelholder 1.' = take('soapbar 1', 'towelholder 1')\n",
    "    def take(self, object, receptacle):\n",
    "        if self.holding is None:\n",
    "            self.holding = object\n",
    "            return self.observation(f'take {object} from {receptacle}')\n",
    "        \n",
    "    # Put an object in or on a receptacle if the agent is holding it. It returns an observation in natural language.\n",
    "    # For example, 'You put the soapbar 1 in/on the cabinet 1.' = put('soapbar 1', 'cabinet 1')\n",
    "    def put(self, object, receptacle):\n",
    "        if self.holding == object:\n",
    "            self.holding = None\n",
    "            return self.observation(f'put {object} in/on {receptacle}')\n",
    "\n",
    "    # Open a receptacle and observe its contents. It returns an observation in natural language.\n",
    "    # For example, 'You open the cabinet 1. The cabinet 1 is open. In it, you see a cloth 1.' = open_receptacle('cabinet 1')\n",
    "    def open_receptacle(self, receptacle):\n",
    "        return self.observation(f'open {receptacle}')\n",
    "\n",
    "    # Close an opened receptacle. It returns an observation in natural language.\n",
    "    # For example, 'You close the safe 1.' = close_receptacle('safe 1')\n",
    "    def close_receptacle(self, receptacle):\n",
    "        return self.observation(f'close {receptacle}')\n",
    "\n",
    "    # Clean an object with a receptacle. It returns an observation in natural language.\n",
    "    # For example, 'You clean the soapbar 1 using the sinkbasin 1.' = clean('soapbar 1', 'sinkbasin 1')\n",
    "    def clean(self, object, receptacle):\n",
    "        return self.observation(f'clean {object} with {receptacle}')\n",
    "\n",
    "    # Heat an object with a receptacle. It returns an observation in natural language.\n",
    "    # For example, 'You heat the tomato 1 using the microwave 1.' = heat('tomato 1', 'microwave 1')\n",
    "    def heat(self, object, receptacle):\n",
    "        return self.observation(f'heat {object} with {receptacle}')\n",
    "\n",
    "    # Cool an object with a receptacle. It returns an observation in natural language.\n",
    "    # For example, 'You cool the pan 2 using the fridge 1.' = cool('pan 2', 'fridge 1')\n",
    "    def cool(self, object, receptacle):\n",
    "        return self.observation(f'cool {object} with {receptacle}')\n",
    "\n",
    "    # Turn_on an object. It returns an observation in natural language.\n",
    "    # For example, 'You turn on the desklamp 1.' = turn_on('desklamp 1')\n",
    "    def turn_on(self, object):\n",
    "        return self.observation(f'use {object}')\n",
    "    \n",
    "    # Report agent's current state, including its location, what it's holding, and last three actions and observations.\n",
    "    # This function should only be used for assertion.\n",
    "    def report(self):\n",
    "        msg = \\\n",
    "f'''The last three interactions before error were:\n",
    "Act: {self.interaction_history[\"actions\"][-3]}\n",
    "Obs: {self.interaction_history[\"observations\"][-3]}\n",
    "Act: {self.interaction_history[\"actions\"][-2]}\n",
    "Obs: {self.interaction_history[\"observations\"][-2]}\n",
    "Act: {self.interaction_history[\"actions\"][-1]}\n",
    "Obs: {self.interaction_history[\"observations\"][-1]}\n",
    "I am at {self.location} and holding {self.holding}.\n",
    "'''.strip()\n",
    "        return msg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Closed-loop simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# dict for storing the failed_tasks_id task numbers w.r.t. each task type\n",
    "\n",
    "failed_tasks_id = {task_name: [] for task_name in prefixes.keys()}\n",
    "num_refinement = {task_name: [] for task_name in prefixes.keys()}\n",
    "all_tasks_id = {task_name: [] for task_name in prefixes.keys()}\n",
    "\n",
    "# adapted from ReAct code.\n",
    "cnts = [0] * 6\n",
    "rs = [0] * 6\n",
    "\n",
    "for _ in range(134):\n",
    "    terminal_output = ''\n",
    "    description, task_name = start_episode(env)\n",
    "    receptacle_list, task = extract_receptacles_and_task(description)\n",
    "\n",
    "    print(task_name)\n",
    "    # define environment \n",
    "    agent = Agent(receptacle_list)\n",
    "\n",
    "    # define the task\n",
    "    print(task)\n",
    "\n",
    "    if task_name.startswith('pick_two_obj'):\n",
    "        example = puttwo_example\n",
    "    elif task_name.startswith('look_at_obj'):\n",
    "        example = examine_example\n",
    "    elif task_name.startswith('pick_and_place'):\n",
    "        example = simple_example\n",
    "    elif task_name.startswith('pick_clean_then_place'):\n",
    "        example = clean_example\n",
    "    elif task_name.startswith('pick_heat_then_place'):\n",
    "        example = heat_example\n",
    "    elif task_name.startswith('pick_cool_then_place'):\n",
    "        example = cool_example\n",
    "    \n",
    "    # get the solution function\n",
    "    prompt = get_solution_prompt\\\n",
    "            .replace('<receptacle_list>', str(receptacle_list))\\\n",
    "            .replace('<task>', task)\\\n",
    "            .replace('<example>', example)\n",
    "    response = ask(prompt)\n",
    "\n",
    "    # refine internally                \n",
    "    solution_func = '''\n",
    "    def solution(agent, start_from=1):\n",
    "        <solution>\n",
    "    '''.strip().replace('<solution>', response) if not response.startswith('def solution(agent, start_from=1):') else response\n",
    "\n",
    "    prompt = code_check_prompt\\\n",
    "            .replace('<solution_func>', solution_func)\n",
    "    response = ask(prompt)\n",
    "    answers = extract_answers(response, ['[Decision]', '[Revised code]'])\n",
    "    # if there is a No after [1]:\n",
    "    if 'Yes' in answers[0]:\n",
    "        print('Fix error in solution function')\n",
    "        solution_func = answers[1].strip('```').replace('Revised code:', '').strip()\n",
    "\n",
    "    # formalize the solution function\n",
    "    solution_func = solution_func.replace('CD', 'cd').replace('solution(agent)','').replace('solution(agent)','').replace('<EOC>', '')\\\n",
    "                                    .replace('print(\"Task completed successfully!\")', '').replace('print(agent.report())', '').replace('receptacles =', '# receptacles =')    \n",
    "    start_num = None\n",
    "    for num_try in range(NUM_TRY):\n",
    "        if num_try < NUM_TRY_RESET:\n",
    "            step = start_num if start_num else 1\n",
    "        else:\n",
    "            step = 1\n",
    "        print('start_from_step:', start_num)\n",
    "        # execute the solution function\n",
    "        def_error = False\n",
    "        try:\n",
    "            exec(solution_func)\n",
    "        except Exception as e:\n",
    "            error_msg = str(e)\n",
    "            error_string = str(e)\n",
    "            checkpoint = None\n",
    "            def_error = True\n",
    "\n",
    "        if not def_error:\n",
    "            error_string, checkpoint, output_string = capture_output(solution, agent, step)\n",
    "            terminal_output += output_string\n",
    "            if error_string:\n",
    "                error_msg = error_string.split('\\n')[4:]\n",
    "                error_msg = '\\n'.join(error_msg)\n",
    "            else:\n",
    "                error_msg = 'You executed the solution function successfully but the task is not completed. Please check your solution function.'\n",
    "        \n",
    "        start_num = None\n",
    "\n",
    "        if agent.is_success:\n",
    "            break\n",
    "\n",
    "        prev_solution_func = solution_func\n",
    "        print(\"Captured error:\", error_string.strip())\n",
    "        error_step = get_error_step(error_string)\n",
    "        # refine based on environment feedback\n",
    "        prompt = feedback_fix_prompt\\\n",
    "                .replace('<example>', example)\\\n",
    "                .replace('<receptacle_list>', str(receptacle_list))\\\n",
    "                .replace('<task>', task)\\\n",
    "                .replace('<error_msg>', error_msg)\n",
    "        response = ask(prompt)\n",
    "        solution_func = '''\n",
    "            def solution(agent, start_from=1):\n",
    "                <solution>\n",
    "            '''.strip().replace('<solution>', response) if not response.startswith('def solution(agent, start_from=1):') else response\n",
    "\n",
    "        # formalize the solution function\n",
    "        solution_func = solution_func.replace('CD', 'cd').replace('solution(agent)','').replace('solution(agent)','').replace('<EOC>', '')\\\n",
    "                                    .replace('print(\"Task completed successfully!\")', '').replace('print(agent.report())', '').replace('receptacles =', '# receptacles =')\n",
    "        \n",
    "        prompt = get_start_from_prompt\\\n",
    "                .replace('<previous_solution>', prev_solution_func)\\\n",
    "                .replace('<revised_solution>', solution_func)\n",
    "        response = ask(prompt)\n",
    "        start_num = get_first_digit(response)\n",
    "        solution_func = solution_func.replace('def solution(agent, start_from=1):', modify_header(checkpoint))\n",
    "\n",
    "\n",
    "    # store results\n",
    "    for i, (k, v) in enumerate(prefixes.items()):\n",
    "        if task_name.startswith(k):\n",
    "            rs[i] += agent.is_success\n",
    "            cnts[i] += 1\n",
    "            num_refinement[k].append(num_try if agent.is_success else -1)\n",
    "            all_tasks_id[k].append(_)\n",
    "            if not agent.is_success:\n",
    "                failed_tasks_id[k].append(_)\n",
    "            break\n",
    "\n",
    "    save_to_file(save_to_folder, LOG_FILE, \n",
    "                f'Task {_+1}: {task_name}\\n' + \\\n",
    "                description + '\\n' + \\\n",
    "                terminal_output + '\\n' + \\\n",
    "                f'Success: {agent.is_success}\\n' + \\\n",
    "                f'Task {_+1}, rs: {rs} cnts {cnts} success_rate: {sum(rs) / sum(cnts)} \\n')\n",
    "    save_to_file(save_to_folder, LOG_FILE, f'failed_tasks_id {failed_tasks_id}')\n",
    "    save_to_file(save_to_folder, LOG_FILE, f'num_refinement {num_refinement}')\n",
    "    save_to_file(save_to_folder, LOG_FILE, f'all_tasks_id {all_tasks_id}')\n",
    "    save_to_file(save_to_folder, LOG_FILE, f'------------\\n')\n",
    "\n",
    "    print(_+1, 'rs', rs, 'cnts', cnts, 'sum(rs)/sum(cnts)', sum(rs) / sum(cnts))\n",
    "    print('------------\\n')\n",
    "print('failed_tasks_id', failed_tasks_id)\n",
    "print('num_refinement', num_refinement)\n",
    "print('all_tasks_id', all_tasks_id)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llms_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "bb48e94b5d9be477c268ed87de4a179163ce583390bda0a928502a1e66012b98"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}